/*
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package net.javlov.example.cont;

import java.util.ArrayList;
import java.util.List;

import net.javlov.Action;
import net.javlov.Agent;
import net.javlov.Environment;
import net.javlov.Option;
import net.javlov.State;
import net.javlov.VectorState;

public class ContinuousWorld implements Environment<double[]> {
	
	/**
	 * Current state the agent is in
	 */
	protected VectorState currentState;
	
	/**
	 * Whether the world wraps around
	 */
	protected boolean wrap;
	
	/**
	 * Width and height
	 */
	protected int width, height;
	
	/**
	 * Goal state, if any
	 */
	protected VectorState goalState;
	
	/**
	 * List of user-defined objects
	 */
	protected List<?> objects;
	
	/**
	 * Goal is reached when within the agent radius
	 */
	public double agentRadius = 1;
	
	protected List<? extends Option<?>> optionPool;
	
	/**
	 * 
	 * @param width
	 * @param height
	 * @param wrap
	 */
	public ContinuousWorld(int width, int height, boolean wrap) {
		this.width = width;
		this.height = height;
		this.wrap = wrap;
		init();
	}
	
	public void setOptions(List<? extends Option<?>> optionPool) {
		this.optionPool = optionPool;
	}
	
	public VectorState getCurrentState() {
		return currentState;
	}

	public boolean isWrap() {
		return wrap;
	}

	public int getWidth() {
		return width;
	}

	public int getHeight() {
		return height;
	}

	public VectorState getGoalState() {
		return goalState;
	}

	@Override
	public double executeAction(Action<double[]> act, Agent<double[]> a) {
		act.execute(a);
		if ( goalReached() )
			currentState.setTerminal(true);
		return -0.1;
	}

	@Override
	public State<double[]> getObservation(Agent<double[]> a) {
		List<Option<?>> elOpts = new ArrayList<Option<?>>(optionPool.size());
		for ( Option o : optionPool )
			if ( o.isEligible(currentState) )
				elOpts.add(o);
		currentState.setOptionSet(elOpts);
		return currentState;
	}

	@Override
	public int getObservationDim() {
		return 3; //x,y,angle
	}

	@Override
	public void init() {
		reset();		
	}

	@Override
	public void reset() {		
		currentState = new VectorState( new double[]{
				agentRadius + Math.random()*(width - 2*agentRadius),
				agentRadius + Math.random()*(height - 2*agentRadius),
				Math.random()*2*Math.PI
		} );
	}
	
	public boolean goalReached() {
		double[] g = goalState.getData(),
				c = currentState.getData();
		return Math.sqrt((c[0]-g[0])*(c[0]-g[0]) + (c[1]-g[1])*(c[1]-g[1])) < agentRadius;
	}
	
	public void setGoal(int x, int y) {
		goalState = new VectorState(new double[]{x,y});
	}
	
	/**
	 * Move distance d into the direction given by the current state.
	 * 
	 * @param d the distance to move
	 * @return true if move succeeded (false if e.g. agent stays in same place at edge of world)
	 */
	public boolean move(double d) {			
		double[] p = currentState.getData(),
				pnew = new double[]{p[0] + d*Math.sin(p[2]), p[1] + d*Math.cos(p[2]), p[2]};
		boolean m = checkWrap(p,pnew);
		currentState = new VectorState(pnew);
		return m;
	}
	
	public boolean turn(double angle) {
		double[] p = currentState.getData(),
				pnew = new double[]{p[0],p[1],(p[2] + angle) % (2*Math.PI)};
		currentState = new VectorState(pnew);
		return true;
	}
	
	protected boolean checkWrap(double[] p, double[] pnew) {
		boolean m = true;
		if ( pnew[0] < agentRadius || pnew[1] < agentRadius 
				|| pnew[0] > width-agentRadius || pnew[1] > height-agentRadius ) {
			if ( wrap ) {
				//move to other side
				pnew[0] = ( pnew[0] < agentRadius ? width-agentRadius : (pnew[0] > width-agentRadius ? agentRadius : pnew[0]) );
				pnew[1] = ( pnew[1] < agentRadius ? height-agentRadius : (pnew[1] > height-agentRadius ? agentRadius : pnew[1]) );
			} else {
				//stay in the same place
				m = false;
				pnew[0] = p[0];
				pnew[1] = p[1];
			}
		}
		return m;
	}
	
	public void setObjects(List<?> objs) {
		objects = objs;
	}
	
	public List<?> getObjects() {
		return objects;
	}

}
